#!/usr/bin/env python
# -*- coding: UTF-8 -*-
'''
@Project ：V2 
@File    ：T2.py
@IDE     ：PyCharm 
@Author  ：郭星
@Date    ：2025/10/8 12:06 
'''
"""
修正版：保证串联设备“同时开始、同时结束”的调度仿真 + GA 求解
- 可配置：瓶数、最大组分数、配气机数、分析设备数
- 支持：可选的 device-specific 分析时长（component_device_time），若未提供则使用 component_time（与设备无关）
- 输出：甘特图（串联用虚线连接）、GA 收敛曲线
"""

import random
import copy
from typing import List, Dict, Tuple
import matplotlib.pyplot as plt
import os

# -------------------- 参数（可修改） --------------------
NUM_BOTTLES = 20          # 瓶数（规模）
MAX_COMPONENTS = 4        # 每瓶最多组分数
NUM_FILLERS = 3           # 配气设备数量（填充机）
NUM_DEVICES = 8           # 分析设备数量
DEVICE_MAX_CAPS = 3       # 每台设备随机能测的最多组件种类（用于生成实例）
POP_SIZE = 100
GENERATIONS = 120
SEED = 123

# 是否使用设备针对组件的不同耗时（True = 使用 component_device_time，如果为 False 则用 component_time）
USE_DEVICE_SPECIFIC_TIMES = False

# -------------------- 随机生成问题实例 --------------------
random.seed(SEED)

# 组件池（按需可扩展）
COMP_POOL = ['A','B','C','D','E','F','G','H'][:MAX_COMPONENTS]

# 为每个瓶随机生成 1..MAX_COMPONENTS 个**有序**组件（不重复）
job_components: Dict[int, List[str]] = {}
for j in range(1, NUM_BOTTLES+1):
    k = random.randint(1, MAX_COMPONENTS)
    # 注意保持顺序：用 sample 会打乱，但保序并不重要，只要表示组分顺序。这里我们用 sample。
    comps = random.sample(COMP_POOL, k)
    job_components[j] = comps

# 每个组件的基础分析时间（与设备无关）——单位：分钟
component_time: Dict[str, float] = {c: round(random.uniform(0.8, 3.5), 2) for c in COMP_POOL}

# 可选：每个组件在每台设备上的耗时（更精细）
# 格式：component_device_time[component][device] = time
component_device_time: Dict[str, Dict[int, float]] = {}
if USE_DEVICE_SPECIFIC_TIMES:
    for c in COMP_POOL:
        component_device_time[c] = {}
        for d in range(1, NUM_DEVICES+1):
            # 随机偏差（基于基础时间）
            component_device_time[c][d] = round(component_time[c] * random.uniform(0.8, 1.4), 2)

# 生成设备能力：每台设备能测的若干组件
device_capability: Dict[int, List[str]] = {}
for d in range(1, NUM_DEVICES+1):
    cap_count = random.randint(1, min(DEVICE_MAX_CAPS, len(COMP_POOL)))
    device_capability[d] = random.sample(COMP_POOL, cap_count)

# 填充机每个组件充气时间（简化：所有组件相同时间，也可扩展成 dict）
filler_time_per_comp = {c: round(random.uniform(0.5, 1.2), 2) for c in COMP_POOL}

# -------------------- 辅助函数 --------------------
def devices_cover_components(devs: List[int], comps: List[str]) -> bool:
    """检查设备集合 devs 能否覆盖 comps 中的所有组件"""
    covered = set()
    for d in devs:
        covered.update(device_capability[d])
    return set(comps).issubset(covered)

def random_partition_sequence(seq: List[str], cut_prob=0.4) -> List[List[str]]:
    """把一条组件序列切成若干连续子组（用于随机生成 plans 的分组）"""
    n = len(seq)
    if n==0:
        return []
    cuts = [False]*(n-1)
    for i in range(n-1):
        cuts[i] = (random.random() < cut_prob)
    groups = []
    cur = [seq[0]]
    for i in range(n-1):
        if cuts[i]:
            groups.append(cur)
            cur = [seq[i+1]]
        else:
            cur.append(seq[i+1])
    groups.append(cur)
    return groups

def random_device_combination_for_group(comps: List[str], max_tries=100) -> List[int]:
    """随机选择若干设备组合以覆盖 comps（贪心随机策略 + 回退穷举）"""
    all_devices = list(device_capability.keys())
    candidates = [d for d in all_devices if any(c in device_capability[d] for c in comps)]
    if not candidates:
        return []
    for _ in range(max_tries):
        chosen = []
        random.shuffle(candidates)
        for d in candidates:
            if random.random() < 0.6 or not chosen:
                chosen.append(d)
            if devices_cover_components(chosen, comps):
                return chosen
    # 小设备数时可以尝试组合穷举以保证能覆盖
    if len(candidates) <= 12:
        from itertools import combinations
        for r in range(1, len(candidates)+1):
            for comb in combinations(candidates, r):
                if devices_cover_components(list(comb), comps):
                    return list(comb)
    return []

# -------------------- 个体生成（order + plans） --------------------
def generate_random_individual(job_ids: List[int]):
    """个体：{'order': [job,...], 'plans': {job: [ {'components': [...], 'devices':[...]} , ... ] }}"""
    order = job_ids[:]
    random.shuffle(order)
    plans = {}
    for j in job_ids:
        comps = job_components[j]
        groups = random_partition_sequence(comps)
        plan_groups = []
        for g in groups:
            devs = random_device_combination_for_group(g)
            if not devs:
                # 退化为单组件分配（找能够测该组件的任意设备）
                for comp in g:
                    possible = [d for d, caps in device_capability.items() if comp in caps]
                    if possible:
                        plan_groups.append({'components':[comp], 'devices':[random.choice(possible)]})
                    else:
                        plan_groups.append({'components':[comp], 'devices':[]})
            else:
                plan_groups.append({'components': g, 'devices': devs})
        plans[j] = plan_groups
    return {'order': order, 'plans': plans}

# -------------------- 关键：计算某组（components + devices）的耗时 --------------------
def compute_group_duration(comps: List[str], devs: List[int]) -> float:
    """
    若有 component_device_time（设备特异性时间），则对每个设备 d 计算
      time_d = max( component_device_time[comp][d] if 有 else component_time[comp] for comp in comps )
    然后组耗时 = max_d time_d
    否则（没有设备特异性）直接组耗时 = max(component_time[comp] for comp in comps)
    解释：设备必须同时开始并同时结束，因此组耗时需能保证每台设备完成其负责的所有组内组件；
    如果设备对组内组件耗时不同，则以设备耗时的最大值为准（所有设备同步结束）。
    """
    if USE_DEVICE_SPECIFIC_TIMES and component_device_time:
        times_per_device = []
        for d in devs:
            t_d = 0.0
            for comp in comps:
                t_comp_on_d = component_device_time.get(comp, {}).get(d, component_time[comp])
                if t_comp_on_d > t_d:
                    t_d = t_comp_on_d
            times_per_device.append(t_d)
        return max(times_per_device) if times_per_device else float('inf')
    else:
        return max(component_time[c] for c in comps)

# -------------------- 仿真：配气 + 分析（保证串联同时开始/结束） --------------------
def simulate_makespan(ind: Dict) -> Tuple[float, Dict]:
    """
    - 配气阶段：有 NUM_FILLERS 台填充机，每台一次只能处理一个瓶的一个组件（组件间在瓶内按顺序）。
      我们让瓶按个体 order 顺序进入配气调度（可扩展成更复杂策略）。
    - 分析阶段：对于每个瓶的每个 group（group 指定了 components 列表和 devices 列表），
      其开始时间 = max( ready_time (fill finished for comps), max(device_free[d] for d in group.devices) )
      结束时间 = start + compute_group_duration(...)
      然后将所有参与设备的 device_free 设为 finish（保证同时结束）
    返回 makespan 和 详细调度信息（用于绘图）
    """
    # 配气模拟
    filler_free = [0.0]*NUM_FILLERS
    fill_finish = {}  # (job, comp) -> finish time
    # 我们按个体中的 order 来做配气（简单、确定）
    for job in ind['order']:
        comps = job_components[job]
        for idx, comp in enumerate(comps):
            prev_finish = 0.0
            if idx > 0:
                prev_finish = fill_finish[(job, comps[idx-1])]
            # 选择能最早完成的 filler
            best_idx = None
            best_finish = None
            for i in range(NUM_FILLERS):
                start_t = max(filler_free[i], prev_finish)
                finish_t = start_t + filler_time_per_comp.get(comp, 1.0)
                if best_finish is None or finish_t < best_finish:
                    best_finish = finish_t
                    best_idx = i
            filler_free[best_idx] = best_finish
            fill_finish[(job, comp)] = best_finish

    # 分析模拟（保证组内设备同步）
    device_free = {d: 0.0 for d in device_capability.keys()}
    group_records = []  # 记录每个 group 的调度，便于甘特图
    # 分析按个体 order
    for job in ind['order']:
        groups = ind['plans'][job]
        for gi, group in enumerate(groups):
            comps = group['components']
            devs = group['devices']
            # 校验：设备集合必须能覆盖组内组件
            if not devs or not devices_cover_components(devs, comps):
                return float('inf'), {'error': f'job {job} group {gi} invalid devices {devs} for comps {comps}'}
            ready_time = max(fill_finish[(job, c)] for c in comps)
            earliest_device_free = max(device_free[d] for d in devs)
            start = max(ready_time, earliest_device_free)   # <-- 关键：所有设备同时开始
            dur = compute_group_duration(comps, devs)
            finish = start + dur
            # 所有设备同时释放
            for d in devs:
                device_free[d] = finish
            group_records.append({'job': job, 'group_idx': gi, 'components': comps, 'devices': devs, 'start': start, 'finish': finish})

    makespan = max((rec['finish'] for rec in group_records), default=0.0)
    detail = {'groups': group_records, 'fill_finish': fill_finish}
    return makespan, detail

# -------------------- GA 操作 --------------------
def order_crossover_OX(p1: List[int], p2: List[int]) -> Tuple[List[int], List[int]]:
    n = len(p1)
    a, b = sorted(random.sample(range(n), 2))
    def ox(x, y):
        child = [None]*n
        child[a:b+1] = x[a:b+1]
        yiter = [v for v in y if v not in child]
        idx = 0
        for i in range(n):
            if child[i] is None:
                child[i] = yiter[idx]; idx+=1
        return child
    return ox(p1,p2), ox(p2,p1)

def plans_crossover(A: Dict[int, List[Dict]], B: Dict[int, List[Dict]]) -> Tuple[Dict[int,List[Dict]], Dict[int,List[Dict]]]:
    a = copy.deepcopy(A); b = copy.deepcopy(B)
    for job in list(A.keys()):
        if random.random() < 0.5:
            a[job], b[job] = b[job], a[job]
    return a, b

def mutate_individual(ind: Dict, p_order=0.2, p_plan=0.4):
    ind = copy.deepcopy(ind)
    # order 变异：swap
    if random.random() < p_order:
        i,j = random.sample(range(len(ind['order'])), 2)
        ind['order'][i], ind['order'][j] = ind['order'][j], ind['order'][i]
    # plans 变异：对若干 job 做重分组或重选设备
    for job in ind['plans'].keys():
        if random.random() < p_plan:
            choice = random.choice(['repartition','redevice','shuffle'])
            if choice == 'repartition':
                comps = job_components[job]
                new_groups = random_partition_sequence(comps)
                new_plan = []
                for g in new_groups:
                    devs = random_device_combination_for_group(g)
                    if not devs:
                        for comp in g:
                            poss = [d for d,caps in device_capability.items() if comp in caps]
                            if poss:
                                new_plan.append({'components':[comp],'devices':[random.choice(poss)]})
                            else:
                                new_plan.append({'components':[comp],'devices':[]})
                    else:
                        new_plan.append({'components':g,'devices':devs})
                ind['plans'][job] = new_plan
            elif choice == 'redevice':
                groups = ind['plans'][job]
                gi = random.randrange(len(groups))
                g = groups[gi]['components']
                devs = random_device_combination_for_group(g)
                if devs:
                    ind['plans'][job][gi]['devices'] = devs
            elif choice == 'shuffle' and len(ind['plans'][job]) >= 2:
                i,j = random.sample(range(len(ind['plans'][job])), 2)
                ind['plans'][job][i], ind['plans'][job][j] = ind['plans'][job][j], ind['plans'][job][i]
    return ind

# -------------------- GA 主流程 --------------------
def run_ga(pop_size=80, generations=100, cx_rate=0.85, mut_rate=0.35):
    job_ids = list(job_components.keys())
    population = [generate_random_individual(job_ids) for _ in range(pop_size)]
    fitnesses = []
    for ind in population:
        f, _ = simulate_makespan(ind)
        fitnesses.append(f)
    best_idx = min(range(len(population)), key=lambda i: fitnesses[i])
    best = population[best_idx]
    best_fit = fitnesses[best_idx]
    history = [best_fit]

    for gen in range(generations):
        new_pop = [best]  # elitism
        while len(new_pop) < pop_size:
            # 锦标赛选择
            parents = random.sample(population, 4)
            p1 = min(parents[:2], key=lambda x: simulate_makespan(x)[0])
            p2 = min(parents[2:], key=lambda x: simulate_makespan(x)[0])
            if random.random() < cx_rate:
                c1_order, c2_order = order_crossover_OX(p1['order'], p2['order'])
                c1_plans, c2_plans = plans_crossover(p1['plans'], p2['plans'])
                child1 = {'order': c1_order, 'plans': c1_plans}
                child2 = {'order': c2_order, 'plans': c2_plans}
            else:
                child1 = copy.deepcopy(p1); child2 = copy.deepcopy(p2)
            if random.random() < mut_rate:
                child1 = mutate_individual(child1)
            if random.random() < mut_rate:
                child2 = mutate_individual(child2)
            new_pop.append(child1)
            if len(new_pop) < pop_size:
                new_pop.append(child2)
        population = new_pop
        fitnesses = [simulate_makespan(ind)[0] for ind in population]
        cur_best_idx = min(range(len(population)), key=lambda i: fitnesses[i])
        if fitnesses[cur_best_idx] < best_fit:
            best_fit = fitnesses[cur_best_idx]
            best = population[cur_best_idx]
        history.append(best_fit)
        if (gen+1) % 10 == 0 or gen==0:
            print(f"Gen {gen+1}: best makespan = {best_fit:.4f}")
    best_makespan, best_detail = simulate_makespan(best)
    return {'best_individual': best, 'makespan': best_makespan, 'detail': best_detail, 'history': history}

# -------------------- 绘图：甘特图（串联用虚线连接） --------------------
def draw_gantt(detail: Dict, filename='gantt.png'):
    groups = detail['groups']
    # 构建设备时间线
    device_timeline = {d: [] for d in device_capability.keys()}
    for rec in groups:
        for d in rec['devices']:
            device_timeline[d].append({'job': rec['job'], 'group': rec['group_idx'], 'start': rec['start'], 'finish': rec['finish'], 'comps': rec['components']})
    for d in device_timeline:
        device_timeline[d].sort(key=lambda x: x['start'])
    # 画图
    fig, ax = plt.subplots(figsize=(12, 6))
    devs_sorted = sorted(device_timeline.keys(), reverse=True)
    y_map = {d: i*10 for i,d in enumerate(devs_sorted)}
    colors = plt.cm.tab20(range(0, len(job_components)))
    for i, d in enumerate(devs_sorted):
        for seg in device_timeline[d]:
            ax.broken_barh([(seg['start'], seg['finish']-seg['start'])], (y_map[d]-3, 6), facecolor=colors[seg['job']%len(colors)], edgecolor='black')
            ax.text(seg['start']+0.01, y_map[d], f"J{seg['job']}G{seg['group']}", va='center', fontsize=8)
    # 串联虚线连接（在每组的 start 时间纵向画短破折线并水平连接）
    for rec in groups:
        if len(rec['devices']) >= 2:
            ylist = [y_map[d] for d in rec['devices']]
            x = rec['start']
            # 竖短线
            for y in ylist:
                ax.plot([x, x], [y-2, y+2], linestyle='--', linewidth=0.8)
            # 水平连接（小段）
            miny, maxy = min(ylist), max(ylist)
            ax.plot([x, x + max(0.2, 0.05*(rec['finish']-rec['start']))], [miny, maxy], linestyle='--', linewidth=0.8)
    ax.set_yticks([y_map[d] for d in devs_sorted])
    ax.set_yticklabels([f"Device {d}" for d in devs_sorted])
    ax.set_xlabel("Time (minutes)")
    ax.set_title("分析设备甘特图（串联设备以虚线连接）")
    plt.tight_layout()
    plt.savefig(filename)
    plt.close(fig)
    print("Saved Gantt to", filename)

def draw_convergence(history, filename='convergence.png'):
    plt.figure(figsize=(6,4))
    plt.plot(history, marker='o')
    plt.xlabel("Generation")
    plt.ylabel("Best makespan")
    plt.title("GA Convergence")
    plt.grid(True)
    plt.tight_layout()
    plt.savefig(filename)
    plt.close()
    print("Saved convergence to", filename)

# -------------------- 运行示例 --------------------
if __name__ == "__main__":
    print("Problem summary:")
    print("Num bottles:", len(job_components))
    print("Sample jobs (first 6):", {k: job_components[k] for k in list(job_components.keys())[:6]})
    print("Component base times:", component_time)
    print("Device capabilities (sample):", {k: device_capability[k] for k in list(device_capability.keys())[:6]})
    res = run_ga(pop_size=POP_SIZE, generations=GENERATIONS, cx_rate=0.85, mut_rate=0.35)
    print("Best makespan:", res['makespan'])
    best = res['best_individual']
    print("Best order (first 10):", best['order'][:10])
    # 保存图像
    out_g = 'gantt.png'
    out_c = 'convergence.png'
    draw_gantt(res['detail'], filename=out_g)
    draw_convergence(res['history'], filename=out_c)
    print("Done. Files:", out_g, out_c)
